import torch
from torch import nn
from net import MyAlexNet
import numpy as np
from torch.optim import lr_scheduler
import os

from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

import matplotlib.pyplot as plt

# 解决中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

ROOT_TRAIN = r'data\train'
ROOT_TEST = r'data\val'

# 将图像的像素值归一化到【-1， 1】之间
# 可以减小数据的差异，使得模型训练更加稳定可靠。
# 目的是将图像的像素值统一缩放到一个固定的范围内，使其更易于易于处理和训练。
normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

# 数据增强和数据预处理
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomVerticalFlip(),
    # 以0.5的概率对图像进行垂直翻转。通过随机翻转，可以扩充数据集，增加模型的稳定性和鲁棒性。
    transforms.ToTensor(),
    # 将 PIL.Image 类型的图像转换为 PyTorch 可以处理的张量类型。转换后，图像数据的像素值归一化到 [0, 1] 范围内。
    normalize])

val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    normalize])

train_dataset = ImageFolder(ROOT_TRAIN, transform=train_transform)
val_dataset = ImageFolder(ROOT_TEST, transform=val_transform)

train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=32, shuffle=True)
# shuffle表示是否打乱数据顺序

device = 'cuda' if torch.cuda.is_available() else 'cpu'

model = MyAlexNet().to(device)

# 定义一个损失函数
loss_fn = nn.CrossEntropyLoss()
# 交叉熵损失函数，用来度量模型的分类任务中预测值与真实标签之间的差异。
# 对于具有多个类别的分类任务，交叉熵损失函数可以通过将每个类别的预测概率与真实标签的独热编码进行比较，来计算预测结果的误差。

# 定义一个优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# SGD随机梯度下降优化器
# 思想是根据当前损失函数值的梯度方向更新模型参数，逐步寻找损失函数的最小点。
# model.parameters()，指定优化器作用的参数目标，即模型对象的参数
# lr 学习率，用来控制模型参数在每次迭代时的更新步长。学习率设置过大会导致模型不收敛，而设置过小则收敛速度过慢
# momentum 动量，用来控制模型参数更新方向的变化。
# 动量的作用是在参数更新时继续更新上一次的方向，使得参数移动更加平稳，避免震荡。通常来讲，动量的值越大，参数更新方向越平稳，模型收敛速度越快。
# 但值过大会产生欠阻尼振荡，值过小会导致不可控的震荡。
# 在实践中，常见的动量取值范围是 [0.5, 0.99]。

# 学习率每隔10轮变为原来的0.5
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
# 为了让模型在训练过程中更好地收敛，并避免出现不必要的震荡，需要在迭代过程中动态地调整学习率。
# StepLR 方法，即每隔 step_size 个 epoch，将学习率乘以 gamma 的值进行调整。
# 这样的学习率调整策略可以逐步降低学习率，使得模型参数收敛得更慢，避免过拟合。
# 在每个 epoch 结束后，可以调用 lr_scheduler.step() 方法自动更新学习率。

# 定义训练函数
# enumerate() 函数将一个可迭代对象转换为一个枚举对象，
# 返回的元素是一个二元组 (index, item)，index 表示元素在可迭代对象中的位置，item 表示元素本身。
# 在这里，使用 enumerate() 函数可以同时获取 batch 索引 batch 和数据批次 (x, y)，便于后续处理。
def train(dataloader, model, loss_fn, optimizer):
    loss, current, n = 0.0, 0.0, 0
    for batch, (x, y) in enumerate(dataloader):
        image, y = x.to(device), y.to(device)
        output = model(image)
        cur_loss = loss_fn(output, y)
        # 将模型输出 output 和真实标签 y 作为参数传入损失函数 loss_fn 中，计算出当前批次的训练误差 cur_loss。
        _, pred = torch.max(output, axis=1)
        # pred是预测数据的标签
        # torch.max() 函数可以返回一个张量中最大值及其索引值，
        # 这里的 output 张量的形状是 (batch_size, num_classes)，
        # 表示模型对输入数据生成的每个类别的得分。
        # 我们可以通过指定 axis=1 参数，将 torch.max() 函数的作用范围限制在第 2 个维度上，
        # 即对每个样本分别找到预测得分最高的类别。
        # 使用 _ 表示不需要返回第一个变量，而 pred 变量表示返回第二个变量，也就是输出张量 output 中预测得分最高的类别索引，即样本预测的标签。
        cur_acc = torch.sum(y==pred) / output.shape[0]
        # output 张量的形状是 (batch_size, num_classes)，
        # 其中 batch_size 表示当前批次的样本数，
        # 由于一个批次中的样本数是相同的，因此可以用 output.shape[0] 来表示当前批次的样本数。

        # 反向传播
        # 其作用是通过计算误差并将误差反向传播至网络中的权重和偏置参数，从而更新参数使得网络能够更好地拟合训练数据。
        optimizer.zero_grad()
        # 将模型参数的梯度清零的函数。
        # 在使用梯度下降算法进行参数优化时，我们通常需要在每轮迭代之前清空参数的梯度，避免上一轮迭代的梯度对本次迭代的结果造成影响。
        cur_loss.backward()
        # 用于计算当前损失的梯度
        # 通过计算当前损失相对于模型参数的梯度，可以使用优化算法（如随机梯度下降）来更新这些参数，从而提高模型的性能。
        # backward()方法则是在当前损失上调用PyTorch的自动微分机制，计算出相对于模型参数的梯度。
        optimizer.step()
        # 通过反向传播计算出模型参数的梯度后，我们可以通过调用 optimizer.step() 来更新这些参数。
        loss += cur_loss.item()
        current += cur_acc.item()
        n = n+1

    train_loss = loss / n
    train_acc = current / n
    print('train_loss' + str(train_loss))
    print('train_acc' + str(train_acc))
    return train_loss, train_acc

# 定义一个验证函数
def val(dataloader, model, loss_fn):
    # 将模型转化为验证模型
    model.eval()
    # 这步是为了防止Dropout在验证集中使用
    loss, current, n = 0.0, 0.0, 0
    with torch.no_grad():
        for batch, (x, y) in enumerate(dataloader):
            image, y = x.to(device), y.to(device)
            output = model(image)
            cur_loss = loss_fn(output, y)
            _, pred = torch.max(output, axis=1)
            cur_acc = torch.sum(y == pred) / output.shape[0]
            loss += cur_loss.item()
            current += cur_acc.item()
            n = n + 1

    val_loss = loss / n
    val_acc = current / n
    print('val_loss' + str(val_loss))
    print('val_acc' + str(val_acc))
    return val_loss, val_acc

# 定义画图函数
def matplot_loss(train_loss, val_loss):
    plt.plot(train_loss, label='train_loss')
    plt.plot(val_loss, label='val_loss')
    plt.legend(loc='best')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.title("训练集和验证集loss值对比图")
    plt.show()

def matplot_acc(train_acc, val_acc):
    plt.plot(train_acc, label='train_acc')
    plt.plot(val_acc, label='val_acc')
    plt.legend(loc='best')
    plt.ylabel('acc')
    plt.xlabel('epoch')
    plt.title("训练集和验证集acc值对比图")
    plt.show()



# 开始训练
loss_train = []
acc_train = []
loss_val = []
acc_val = []

epoch = 20
min_acc = 0
for t in range(epoch):
    lr_scheduler.step()
    # 自动更新学习率
    print(f"epoch{t+1}\n-----------")
    train_loss, train_acc = train(train_dataloader, model, loss_fn, optimizer)
    val_loss, val_acc = val(val_dataloader, model, loss_fn)

    loss_train.append(train_loss)
    acc_train.append(train_acc)
    loss_val.append(val_loss)
    acc_val.append(val_acc)

    # 保存最好的模型权重
    if val_acc >min_acc:
        folder = 'save_model'
        if not os.path.exists(folder):
            os.mkdir('save_model')
        min_acc = val_acc
        print(f"save best model, 第{t+1}轮")
        torch.save(model.state_dict(), 'save_model/best_model.pth')
    # 保存最后一轮的权重文件
    if t == epoch-1:
        torch.save(model.state_dict(), 'save_model/last_model.pth')

matplot_loss(loss_train, loss_val)
matplot_acc(acc_train, acc_val)
print('Done!')
